#! /usr/bin/env python
import rospy
from turtlesim.msg import Pose
from geometry_msgs.msg import Twist
from enum import Enum
import math

long = 3    # 长边长
short = 2   # 短边长
g_pose = Pose()     # 当前pose
g_goal = Pose()     # 目标pose

# 枚举定义小乌龟的几种状态
class State(Enum):
  FORWARD=1
  STOP_FORWARD=2
  TURN=3
  STOP_TURN=4

g_state = State.FORWARD     # 小乌龟当前状态，初始化为FORWARD
g_last_state = State.FORWARD    # 小乌龟的上一个状态，初始化为FORWARD
g_first_goal_set = False        # 是否已经设置了初始目标？否
  
# 接收到小乌龟Pose消息时的回调函数
def poseCallback(msg):
  global g_pose
  g_pose = msg
  # print(g_pose) 

# 是否已经到达目标？
def hasReachedGoal():
  return abs(g_pose.x - g_goal.x)<0.1 and abs(g_pose.y-g_goal.y)<0.1 and abs(g_pose.theta-g_goal.theta)<0.01

# 是否已经停止运动？
def hasStopped():
  return g_pose.angular_velocity<0.0001 and g_pose.linear_velocity<0.0001

def printGoal():
  print("New goal [%f, %f, %f]" % (g_goal.x, g_goal.y, g_goal.theta))

def commandTurtle(twist_pub, linear, angular):
  twist = Twist()
  twist.linear.x = linear
  twist.angular.z = angular
  twist_pub.publish(twist)

def stopForward(twist_pub):
  global g_state
  global g_goal
  if hasStopped():
    print("Reached goal")
    g_state =  State.TURN
    g_goal.x = g_pose.x
    g_goal.y = g_pose.y
    g_goal.theta = (g_pose.theta + math.pi/2.0) % (2*math.pi)
    if (g_goal.theta>=math.pi):
      g_goal.theta -= 2*math.pi           # 将角度转换为在[-pi, pi]
    printGoal()
  else:
    commandTurtle(twist_pub, 0, 0)

def stopTurn(twist_pub):
  global g_state
  global g_goal
  if hasStopped():
    print("Reached goal")
    g_state = State.FORWARD
    g_goal.x = math.cos(g_pose.theta)*long + g_pose.x
    g_goal.y = math.sin(g_pose.theta)*short + g_pose.y
    g_goal.theta = g_pose.theta
    printGoal()
  else:
    commandTurtle(twist_pub, 0, 0)

def forward(twist_pub):
  global g_state
  if hasReachedGoal():
    g_state = State.STOP_FORWARD
    commandTurtle(twist_pub, 0, 0)
  else:
    commandTurtle(twist_pub, 1.0, 0.0)

def turn(twist_pub):
  global g_state
  if hasReachedGoal():
    g_state = State.STOP_TURN
    commandTurtle(twist_pub, 0, 0)
  else:
    commandTurtle(twist_pub, 0.0, 0.4)

#======== 主程序 ========
if __name__ == '__main__':
  twist_pub = rospy.Publisher('/turtle1/cmd_vel', Twist, queue_size=1)
  pos_sub = rospy.Subscriber('/turtle1/pose', Pose,poseCallback) 
  rospy.init_node('turtle_square')
  rate = rospy.Rate(100) 

  while not rospy.is_shutdown():
    if (g_first_goal_set==False):    # 设置初始目标
      g_first_goal_set = True
      g_state = State.FORWARD
      g_goal.x = math.cos(g_pose.theta)*long + g_pose.x 
      g_goal.y = math.sin(g_pose.theta)*short +g_pose.y
      g_goal.theta = g_pose.theta
      printGoal()
    
    if (g_state == State.FORWARD):
      forward(twist_pub)
    elif (g_state==State.STOP_FORWARD):
      stopForward(twist_pub)
    elif (g_state==State.TURN):
      turn(twist_pub)
    elif (g_state==State.STOP_TURN):
      stopTurn(twist_pub)

    rate.sleep()






